#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define TOTAL_MEM_SIZE 8096
#define MIN_BLOCK_SIZE 16

typedef struct block {
    int size;
    struct block *next;
    struct block *prev;
} Block;

void split(Block *b);
Block *find_best_fit(int size);
void print_blocks();
void print_block(Block *b);

char mem[TOTAL_MEM_SIZE];
int num_levels = 0; // 计算需要多少个块大小等级
Block *free_lists[32] = {NULL}; // 空闲块链表

void init() {
    // 计算需要多少个块大小等级
    int i = 1;
    while (MIN_BLOCK_SIZE << i <= TOTAL_MEM_SIZE) {
        i++;
    }
    num_levels = i;

    // 初始化最大块
    Block *max_block = (Block *)mem;
    max_block->size = TOTAL_MEM_SIZE;
    max_block->next = NULL;
    max_block->prev = NULL;

    // 将最大块加入空闲链表
    free_lists[num_levels - 1] = max_block;
}

void *umalloc(size_t size) {
    if (size == 0) {
        return NULL;
    }

    // 计算需要的块大小等级
    int level = 0;
    int block_size = MIN_BLOCK_SIZE;
    while (block_size < size + sizeof(Block)) {
        block_size <<= 1;
        level++;
    }

    // 查找合适的空闲块
    Block *b = find_best_fit(block_size);
    if (b == NULL) {
        return NULL;
    }

    // 如果找到的块比需要的块大，就分裂成更小的块
    while (b->size > block_size + sizeof(Block)) {
        split(b);
        b = b->next;
    }

    // 从空闲链表中移除该块
    if (b->prev != NULL) {
        b->prev->next = b->next;
    }
    else {
        free_lists[level] = b->next;
    }
    if (b->next != NULL) {
        b->next->prev = b->prev;
    }

    // 返回可用内存的指针
    return (void *)(b + 1);
}

void ufree(void *ptr) {
    if (ptr == NULL) {
        return;
    }

    // 获取块地址
    Block *b = (Block *)ptr - 1;

    // 加入空闲链表
    Block *current_block = b;
    while (current_block->size < TOTAL_MEM_SIZE) {
        // 计算伙伴块地址
        int buddy_offset = ((char *)current_block - mem) ^ current_block->size;
        Block *buddy = (Block *)(mem + buddy_offset);

        // 判断伙伴块是否空闲并且大小相同
        if (buddy->size == current_block->size && (char *)buddy >= mem && (char *)buddy < mem + TOTAL_MEM_SIZE) {
            // 从空闲链表中移除伙伴块
            if (buddy->prev != NULL) {
                buddy->prev->next = buddy->next;
            }
            else {
                free_lists[num_levels - 1] = buddy->next;
            }
            if (buddy->next != NULL) {
                buddy->next->prev = buddy->prev;
            }

            // 将两个块合并成一个更大的块
            current_block = current_block < buddy ? current_block : buddy;
            current_block->size *= 2;
            current_block->next = NULL;
            current_block->prev = NULL;
        }
        else {
            break;
        }
    }

    // 将合并后的块加入空闲链表
    int level = 0;
    int block_size = MIN_BLOCK_SIZE;
    while (block_size < current_block->size) {
        block_size <<= 1;
        level++;
    }
    current_block->prev = NULL;
    current_block->next = free_lists[level];
    if (free_lists[level] != NULL) {
        free_lists[level]->prev = current_block;
    }
    free_lists[level] = current_block;
}

void umerge() {
    // 遍历所有块，找出需要合并的块
    Block *current_block = (Block *)mem;
    while (current_block->size < TOTAL_MEM_SIZE) {
        // 计算伙伴块地址
        int buddy_offset = ((char *)current_block - mem) ^ current_block->size;
        Block *buddy = (Block *)(mem + buddy_offset);

        // 判断伙伴块是否空闲并且大小相同
        if (buddy->size == current_block->size && (char *)buddy >= mem && (char *)buddy < mem + TOTAL_MEM_SIZE) {
            // 从空闲链表中移除伙伴块
            if (buddy->prev != NULL) {
                buddy->prev->next = buddy->next;
            }
            else {
                int level = 0;
                int block_size = MIN_BLOCK_SIZE;
                while (block_size < buddy->size) {
                    block_size <<= 1;
                    level++;
                }
                free_lists[level] = buddy->next;
            }
            if (buddy->next != NULL) {
                buddy->next->prev = buddy->prev;
            }

            // 将两个块合并成一个更大的块
            current_block = current_block < buddy ? current_block : buddy;
            current_block->size *= 2;
            current_block->next = NULL;
            current_block->prev = NULL;
        }
        else {
            current_block = current_block->next;
        }
    }
}

void split(Block *b) {
    // 计算分裂后块的大小和地址
    int new_size = b->size / 2;
    Block *new_b = (Block *)((char *)b + new_size);

    // 设置分裂后块的属性
    new_b->size = new_size;
    new_b->next = b->next;
    new_b->prev = b;

    // 将分裂后块加入空闲链表
    int level = 0;
    int block_size = MIN_BLOCK_SIZE;
    while (block_size < new_size) {
        block_size <<= 1;
        level++;
    }
    new_b->next = free_lists[level];
    if (free_lists[level] != NULL) {
        free_lists[level]->prev = new_b;
    }
    free_lists[level] = new_b;

    // 更新原块的属性
    b->size = new_size;
    b->next = new_b;
}

Block *find_best_fit(int size) {
    // 查找最小的合适块
    int level = 0;
    int block_size = MIN_BLOCK_SIZE;
    while (block_size < size) {
        block_size <<= 1;
        level++;
    }
    Block *current_block = free_lists[level];
    while (current_block != NULL) {
        if (current_block->size >= size) {
            return current_block;
        }
        current_block = current_block->next;
    }

    // 如果找不到合适块，则返回 NULL
    return NULL;
}

void print_blocks() {
    printf("Memory blocks:\n");
    Block *current_block = (Block *)mem;
    while (current_block->size < TOTAL_MEM_SIZE) {
        print_block(current_block);
        current_block = current_block->next;
    }
    print_block(current_block);
}

void print_block(Block *b) {
    printf("Block %d bytes at %p\n", b->size, b);
}

int main() {
    init();

    void *p1 = umalloc(16);
    void *p2 = umalloc(32);
    void *p3 = umalloc(64);
    void *p4 = umalloc(128);

    print_blocks();

    ufree(p2);

    print_blocks();

    umerge();

    print_blocks();

    return 0;
}
